import time

import faiss
import numpy

numpy.random.seed(0) #设置随机数种子，保证每次运行结果一致
vectors = numpy.random.rand(1000000, 768)  # 创建1000000个随机向量 , 768表示向量的维度
query = numpy.random.rand(1, 768) # 创建1个随机向量 ,用来模拟查询向量

# 使用普通线性索引
def fun_1():
    index = faiss.IndexFlatL2(768)
    index.add(vectors)

    start = time.time()
    D , I = index.search(query , k=2)
    print(f"普通线性索引耗时：{time.time() - start}")
    print(f"最相似的向量索引：{I}")
    print(f"最相似的向量距离：{D}")


#使用IVF索引
def fun_2():
    quantizer = faiss.IndexFlatL2(768) # 创建一个 quantizer, 用于将向量进行量化,也就是对每个分类里的数据查询方式
    index = faiss.IndexIVFFlat(quantizer , 768 , 100) #100 表示每个分类里的数据数量
    #index.nprobe = 2 # 搜索的分类数量,在最相似的N个类中搜索最相似的向量，默认为1
    index.train(vectors) #训练IVF索引,得到100个分类

    index.add(vectors)
    start = time.time()
    D, I = index.search(query, k=2)
    print(f"IVF索引耗时：{time.time() - start}")
    print(f"最相似的向量索引：{I}")
    print(f"最相似的向量距离：{D}")


if __name__ == '__main__':
    fun_1()
    print("---------------------------------------------")
    fun_2()